import spacy
import coreferee
from spacy.kb import KnowledgeBase

# --- 模拟一个预先构建好的、高质量的四川景点知识库 ---
# 在真实应用中，这些信息来源于你离线处理后的结果
KNOWLEDGE_BASE = {
    "SC_JZG": {
        "name": "九寨沟",
        "description": "九寨沟是一个美丽的自然保护区，以其多彩的湖泊、瀑布和雪山而闻名。",
        "activities": ["看五彩池", "观赏诺日朗瀑布", "徒步旅行"]
    },
    "SC_LSDF": {
        "name": "乐山大佛",
        "description": "乐山大佛是一尊巨大的石刻弥勒佛坐像，是世界上最大的石佛。",
        "activities": ["乘船观赏全景", "近距离参观佛像", "游览凌云寺"]
    },
    "SC_CSDXM": {
        "name": "成都大熊猫繁育研究基地",
        "description": "这里是观赏和了解中国国宝大熊猫的最佳地点之一。",
        "activities": ["看成年大熊猫", "看熊猫幼崽", "参观熊猫博物馆"]
    }
}


class SichuanRecommender:
    def __init__(self, kb_data):
        print("--- 正在初始化推荐系统 ---")
        # 1. 加载 spaCy 模型和 coreferee
        self.nlp = spacy.load("en_core_web_lg")
        if 'coreferee' not in self.nlp.pipe_names:
            self.nlp.add_pipe('coreferee')

        # 2. 创建并加载用于实体消歧的知识库 (KB)
        self.kb = self._create_kb(kb_data)
        if "entity_linker" not in self.nlp.pipe_names:
            self.nlp.add_pipe("entity_linker", config={"kb": self.kb})
        print("--- 系统已就绪 ---")

    def _create_kb(self, kb_data):
        """根据我们的数据创建一个 spaCy 知识库"""
        kb = KnowledgeBase(vocab=self.nlp.vocab, entity_vector_length=300)
        for entity_id, data in kb_data.items():
            # 使用描述来生成实体的向量
            entity_vector = self.nlp(data['description']).vector
            kb.add_entity(entity=entity_id, freq=100, entity_vector=entity_vector)
            # 将景点的名字作为别名
            kb.add_alias(alias=data['name'], entities=[entity_id], probabilities=[1.0])
        return kb

    def recommend(self, query, conversation_history=""):
        """
        接收用户查询和对话历史，返回推荐信息
        """
        # 步骤一：使用共指消解来理解上下文
        # 将历史和当前查询拼接，以解析代词
        full_text = conversation_history + " " + query
        doc = self.nlp(full_text)

        # 获取当前查询中被解析的指代
        # 例如，如果 query 是 "它有什么好玩的？"，而历史是"我想了解九寨沟"，
        # resolved_query 会变成 "九寨沟有什么好玩的？"
        resolved_query = query
        last_turn_doc = doc[len(conversation_history) + 1:]  # 只看当前轮次的token
        for token in last_turn_doc:
            if token._.in_coref:
                main_ref = token._.coref_chains[0].main
                # 如果代词的主指代在对话历史中，我们用它来替换
                if main_ref.start < len(conversation_history):
                    resolved_query = resolved_query.replace(token.text, doc[main_ref.start:main_ref.end].text)

        print(f"原始查询: '{query}'")
        print(f"共指消解后查询: '{resolved_query}'")

        # 步骤二：使用实体消歧来识别具体景点
        resolved_doc = self.nlp(resolved_query)
        found_attractions = []
        for ent in resolved_doc.ents:
            if ent.kb_id_:  # 如果实体成功链接到了我们的知识库
                found_attractions.append(ent.kb_id_)

        # 步骤三：根据识别出的景点ID生成回复
        if not found_attractions:
            return "抱歉，我不太明白您指的是哪个景点。可以提供更具体的名字吗？"

        # 只处理第一个识别出的景点
        attraction_id = found_attractions[0]
        attraction_data = KNOWLEDGE_BASE.get(attraction_id)

        if attraction_data:
            # 这是一个简单的回复生成逻辑
            response = (
                f"好的，关于【{attraction_data['name']}】:\n"
                f"简介: {attraction_data['description']}\n"
                f"在这里您可以: {', '.join(attraction_data['activities'])}。"
            )
            return response
        else:
            return "抱歉，知识库中没有这个景点的详细信息。"


# --- 主程序 ---
if __name__ == "__main__":
    recommender = SichuanRecommender(KNOWLEDGE_BASE)

    # 模拟多轮对话
    print("\n--- 对话场景 ---")

    # 第一轮
    history = ""
    query1 = "给我介绍一下成都大熊猫基地"
    response1 = recommender.recommend(query1, history)
    print(f"用户: {query1}")
    print(f"系统: {response1}")

    print("\n-----------------\n")

    # 第二轮：用户使用了代词 "那里"
    # 更新对话历史
    history += query1 + ". " + "好的，关于【成都大熊猫繁育研究基地】:..."  # 简化版的系统回复
    query2 = "那里有什么好玩的？"
    response2 = recommender.recommend(query2, history)
    print(f"用户: {query2}")
    print(f"系统: {response2}")